import pygame
import numpy as np
from collections import deque
from numba import njit

# 初始化Pygame
pygame.init()

# 物理参数
G = 4 * np.pi**2
dt = 0.0005
SOFTENING = 1e-6
BASE_SCALE = 150
scale = BASE_SCALE
offset = np.array([0.0, 0.0])
dragging = False

# 屏幕参数
WIDTH, HEIGHT = 1280, 720
CENTER = np.array([WIDTH//2, HEIGHT//2])
screen = pygame.display.set_mode((WIDTH, HEIGHT), pygame.RESIZABLE)
pygame.display.set_caption("天体运动模拟")

# 颜色定义
COLORS = {
    'sun': (255, 255, 0),
    'earth': (100, 149, 237),
    'random': lambda: tuple(np.random.randint(50, 255, 3))
}

# Numba加速函数
@njit
def compute_accel_numba(pos_self, others_pos, others_mass, softening_sq):
    acc = np.zeros(2)
    for i in range(others_pos.shape[0]):
        dx = pos_self[0] - others_pos[i, 0]
        dy = pos_self[1] - others_pos[i, 1]
        distance_sq = dx*dx + dy*dy
        force_dir = np.array([dx, dy]) / np.sqrt(distance_sq + SOFTENING**2)
        acc += -G * others_mass[i] * force_dir / (distance_sq + softening_sq)**1.5
    return acc

@njit
def find_collisions_numba(positions, radii):
    n = positions.shape[0]
    collisions = []
    for i in range(n):
        for j in range(i+1, n):
            dx = positions[i,0] - positions[j,0]
            dy = positions[i,1] - positions[j,1]
            distance_sq = dx*dx + dy*dy
            if distance_sq < (radii[i] + radii[j])**2:
                collisions.append((i, j))
    return collisions

# 坐标转换函数
def to_screen(pos):
    x = pos[0] * scale + CENTER[0] + offset[0]
    y = CENTER[1] - pos[1] * scale + offset[1]
    return (int(x), int(y))

def to_world(pos):
    dx = (pos[0] - CENTER[0] - offset[0]) / scale
    dy = (CENTER[1] - pos[1] + offset[1]) / scale
    return np.array([dx, dy])

# 天体类
class CelestialBody:
    def __init__(self, pos, vel, mass, color, radius_phys, name="Planet", max_orbit_length=500):
        self.name = name
        self.pos = np.array(pos, dtype='float64')
        self.vel = np.array(vel, dtype='float64')
        self.mass = mass
        self.color = color
        self.radius_phys = radius_phys
        self.orbit = deque(maxlen=max_orbit_length)
        self.merge_pending = False
        self.acceleration = np.zeros(2)

    def compute_acceleration(self, others_pos, others_mass, softening_sq):
        self.acceleration = compute_accel_numba(self.pos, others_pos, others_mass, softening_sq)

    def update_position(self, dt):
        self.pos += self.vel * dt + 0.5 * self.acceleration * dt**2

    def update_velocity(self, new_acc, dt):
        self.vel += 0.5 * (self.acceleration + new_acc) * dt
        self.acceleration = new_acc

    def update_orbit(self):
        self.orbit.append(self.pos.copy())

    def merge(self, other):
        total_mass = self.mass + other.mass
        new_color = (
            (self.mass * self.color[0] + other.mass * other.color[0]) / total_mass,
            (self.mass * self.color[1] + other.mass * other.color[1]) / total_mass,
            (self.mass * self.color[2] + other.mass * other.color[2]) / total_mass
        )
        self.color = tuple(map(int, new_color))
        self.pos = (self.mass * self.pos + other.mass * other.pos) / total_mass
        self.vel = (self.mass * self.vel + other.mass * other.vel) / total_mass
        self.mass = total_mass
        self.radius_phys = np.cbrt(self.radius_phys**3 + other.radius_phys**3)
        other.merge_pending = True

    def draw(self, screen):
        pos = to_screen(self.pos)
        radius_screen = max(2, int(self.radius_phys * scale))
        
        # 确保颜色有效
        try:
            pygame.draw.circle(screen, self.color, pos, radius_screen)
        except:
            pygame.draw.circle(screen, (255,255,255), pos, radius_screen)
        
        # 绘制轨道
        if len(self.orbit) > 2:
            points = []
            for p in self.orbit:
                try:
                    points.append(to_screen(p))
                except:
                    continue
            if len(points) > 1:
                pygame.draw.lines(screen, self.color, False, points, 1)

# 初始化天体
bodies = [
    CelestialBody([0, 0], [0, 0], 1.0, COLORS['sun'], 0.1, "Sun", max_orbit_length=300),
    CelestialBody([1, 0], [0, 6.28], 3e-6, COLORS['earth'], 0.01, "Earth", max_orbit_length=1000)
]

# 交互参数
show_help = True
adding_body = False
paused = False
drag_anchor = None
drag_offset_base = np.array([0.0, 0.0])

# 主循环
running = True
clock = pygame.time.Clock()

while running:
    clock.tick(60)
    screen.fill((0, 0, 0))
    
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.VIDEORESIZE:
            WIDTH, HEIGHT = event.size
            CENTER[:] = WIDTH//2, HEIGHT//2
            screen = pygame.display.set_mode((WIDTH, HEIGHT), pygame.RESIZABLE)
        elif event.type == pygame.MOUSEBUTTONDOWN:
            if event.button == 1:
                dragging = True
                drag_anchor = pygame.mouse.get_pos()
                drag_offset_base = offset.copy()
            elif event.button == 3:
                adding_body = True
                start_pos = pygame.mouse.get_pos()
        elif event.type == pygame.MOUSEBUTTONUP:
            if event.button == 1:
                dragging = False
                drag_anchor = None
            elif event.button == 3 and adding_body:
                end_pos = pygame.mouse.get_pos()
                dx = (end_pos[0] - start_pos[0])/scale * 5
                dy = -(end_pos[1] - start_pos[1])/scale * 5
                pos = to_world(start_pos)
                color = COLORS['random']()
                new_body = CelestialBody(pos, [dx, dy], 1e-7, color, 0.00002, "Custom", max_orbit_length=500)
                bodies.append(new_body)
                adding_body = False
        elif event.type == pygame.MOUSEWHEEL:
            if event.y != 0:
                mouse_pos = pygame.mouse.get_pos()
                world_pos = to_world(mouse_pos)
                
                prev_scale = scale
                scale_factor = 1.1 if event.y > 0 else 0.9
                scale *= scale_factor
                scale = np.clip(scale, 10, 5000)
                
                offset[0] = mouse_pos[0] - CENTER[0] - world_pos[0] * scale
                offset[1] = mouse_pos[1] - CENTER[1] + world_pos[1] * scale
                
                offset = np.clip(offset, -1e4, 1e4)
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_h:
                show_help = not show_help
            elif event.key == pygame.K_ESCAPE:
                running = False
            elif event.key == pygame.K_SPACE:
                paused = not paused
            elif event.key == pygame.K_PLUS or event.key == pygame.K_EQUALS:
                dt = min(dt * 1.1, 0.1)
            elif event.key == pygame.K_MINUS:
                dt = max(dt / 1.1, 0.00001)
            elif event.key == pygame.K_UP:
                for body in bodies:
                    body.orbit = deque(body.orbit, maxlen=body.orbit.maxlen + 100)
            elif event.key == pygame.K_DOWN:
                for body in bodies:
                    new_max = max(100, body.orbit.maxlen - 100)
                    body.orbit = deque(body.orbit, maxlen=new_max)

    if dragging and drag_anchor:
        current_pos = pygame.mouse.get_pos()
        dx = current_pos[0] - drag_anchor[0]
        dy = current_pos[1] - drag_anchor[1]
        offset = drag_offset_base + np.array([dx, dy])

    if adding_body:
        pygame.draw.line(screen, (255, 0, 0), start_pos, pygame.mouse.get_pos(), 2)

    if not paused:
        # 处理合并
        active_bodies = [b for b in bodies if not b.merge_pending]
        if len(active_bodies) >= 2:
            np_positions = np.array([b.pos for b in active_bodies])
            np_radii = np.array([b.radius_phys for b in active_bodies])
            merge_pairs = find_collisions_numba(np_positions, np_radii)
            
            for i, j in merge_pairs:
                if active_bodies[i].merge_pending or active_bodies[j].merge_pending:
                    continue
                active_bodies[i].merge(active_bodies[j])
            
            active_bodies = [b for b in active_bodies if not b.merge_pending]
            bodies = active_bodies

        # 物理计算
        if len(active_bodies) > 0:
            # 计算加速度
            others_pos = np.array([b.pos for b in active_bodies])
            others_mass = np.array([b.mass for b in active_bodies])
            softening_sq = SOFTENING**2
            
            old_accelerations = []
            for i, body in enumerate(active_bodies):
                mask = np.ones(len(active_bodies), dtype=bool)
                mask[i] = False
                body.compute_acceleration(others_pos[mask], others_mass[mask], softening_sq)
                old_accelerations.append(body.acceleration.copy())
            
            # 更新位置
            for body in active_bodies:
                body.update_position(dt)
            
            # 计算新加速
            for i, body in enumerate(active_bodies):
                mask = np.ones(len(active_bodies), dtype=bool)
                mask[i] = False
                body.compute_acceleration(others_pos[mask], others_mass[mask], softening_sq)
            
            # 更新速度
            for body, old_acc in zip(active_bodies, old_accelerations):
                body.update_velocity(body.acceleration, dt)
            
            # 更新轨道
            for body in active_bodies:
                body.update_orbit()

    # 绘制所有天体
    for body in bodies:
        body.draw(screen)

    # 显示帮助信息
    if show_help:
        font = pygame.font.SysFont('Arial', 14)
        help_text = [
            "H: Toggle help",
            "Space: Pause/Resume",
            "+/-: Adjust time step",
            "Mouse wheel: Zoom",
            "Right click: Add body"
            f"Scale: {scale:.0f} AU/pixel",
            f"Time step: {dt:.2e} years",
            f"Bodies: {len(bodies)}",
        ]
        for i, line in enumerate(help_text):
            text = font.render(line, True, (255, 255, 255))
            screen.blit(text, (10, 10 + 20*i))

    pygame.display.update()

pygame.quit()